import configparser
from typing import Dict, List, Optional
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
    MilvusClient
)

from ai.llm.param_builder import MilvusParam
import sys


class DatabaseHelper:

    def searchEmbedding(self, q_embeddings, db_param: MilvusParam):
        """
        向量相似度查询
        执行sdk的search方法
        :param q_embeddings: 查询问题的向量
        :param db_param: 数据库查询参数
            anns_field: 指定查询的字段
            param: 操作参数，可以设置的有：
                metric_type: 应用于此操作的度量类型。这应该与为上面指定的向量字段建立索引时使用的方法相同。可能的值是L2、IP和COSINE。
                params: 附加参数
                    offset： 偏移量，用于分页；
                    radius： 检索范围，最小相似性阈值，当metric_type设置为L2时，请确保此值大于range_filter的值。否则，此值应低于range_filter的值。
                    range：  filter-检索范围，将搜索细化到特定相似性范围内的向量。metric_type设置为IP或COSINE时，请确保此值大于radius。否则，该值应低于radius
                    nprobe： 要查询的单位数
                            怎么理解，认为是聚类数量，随着目标输入向量（nq）和要搜索的聚类数量（nprobe）的增加，查询时间急剧增加；聚类最多越准确
            limit: 要返回的实体总数
            expr:  标量过滤条件
            partition_names: 分区名称列表,默认None,如果指定，则查询中只涉及指定的分区。
            output_fields: 返回的字段列表，如果未指定，则仅包括主字段
            timeout: 超时时间
            round_decimal: Milvus将计算出的距离四舍五入到的小数位数。该值默认为-1，表示Milvus跳过计算距离的四舍五入并返回原始值。
        :return:
        """
        pass


class MilvusDatabase(DatabaseHelper):

    def __init__(self):
        config = configparser.ConfigParser()
        classpath = sys.path
        # 如何获取项目根目录
        config.read(classpath[00] + "/ai/ai-config.ini")
        self.host = config.get('milvus', 'host')
        self.port = config.get('milvus', 'port')
        self.token = config.get('milvus', 'token')
        self.db = config.get('milvus', 'db')
        # 与数据库建立连接
        connections.connect(alias="frank", host=self.host, port=self.port,
                            token=self.token, db_name=self.db)

    def searchEmbedding(self, q_embeddings, db_param: MilvusParam):

        db = Collection(name=db_param.collection_name, using="frank")
        return db.search(q_embeddings, db_param.anns_field, db_param.param, db_param.limit, db_param.expr,
                              db_param.partition_names, db_param.output_fields, db_param.timeout,
                              db_param.round_decimal)

# Singleton Mode in Python
milvusDatabase = MilvusDatabase()


def initUserCollection():
    """
    初始化测试表user
    :return:
    """
    # 加载文本
    from langchain.text_splitter import MarkdownHeaderTextSplitter
    with open('file/user.md') as f:
        data = f.read()
    headers_to_split_on = [
        ("#", "Header 1"),
        ("##", "Header 2"),
        ("###", "Header 3")
    ]
    markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
    md_header_splits = markdown_splitter.split_text(data)
    # 文本解析,构建user
    user_list = list()
    desc_list = list()
    for doc in md_header_splits:
        country = doc.metadata.get("Header 1")
        name = doc.metadata.get("Header 2")
        title = doc.metadata.get("Header 3")
        desc = doc.page_content
        desc_list.append(desc)
        user_list.append({
            "name": f'{country}_{name}_{title}',
            "desc": desc
        })
    # 批量向量化
    from llm_chat_client import qianfanLLM
    embedding_list = qianfanLLM.embed_doc(desc_list)
    # 实体完善
    for i, e in enumerate(user_list):
        user_list[i]['embeddings'] = embedding_list[i]
    # 表数据重置
    db = Collection(name="user", using="frank")
    db.delete("id > 1")
    # 数据插入
    db.insert(user_list)


if __name__ == '__main__':
    db = milvusDatabase.db
    # 问题向量化
    query = "守成皇帝是谁？"
    from llm_chat_client import qianfanLLM
    # embeddings = qianfanLLM.embed_query(query)
    embeddings = [0.0340130515396595, 0.07832217961549759, 0.027131065726280212, 0.1215016171336174, -0.03975407034158707, 0.010455867275595665, 0.015734246000647545, -0.09016700088977814, -0.04898872599005699, -0.055512115359306335, 0.07364185899496078, -0.04075038433074951, 0.03798053786158562, 0.021135348826646805, -0.05866818502545357, -0.01364555861800909, -0.0746140256524086, 0.02789485827088356, -0.032312534749507904, -0.010561644099652767, 0.10903181880712509, -0.055083878338336945, -0.0159218180924654, 0.04380373656749725, -0.05123565346002579, -0.1384463608264923, -0.09217094630002975, -0.016642523929476738, 0.03431130200624466, 0.06640094518661499, 0.09185709804296494, -0.0619460865855217, -0.004553212318569422, 0.01095954142510891, 0.10394428670406342, 0.04616805538535118, -0.013809160329401493, -0.023472243919968605, -0.06361102312803268, -0.09887931495904922, -0.044074248522520065, -0.06539227813482285, -0.025161683559417725, -0.0009903713362291455, 0.06380517035722733, -0.017219340428709984, 0.018179025501012802, 0.02797204814851284, 0.054495953023433685, -0.03653597831726074, -0.01600138284265995, -0.09115820378065109, -0.023935239762067795, -0.05858219787478447, 0.12812934815883636, -0.034655604511499405, 0.02040186896920204, 0.03583560138940811, -0.044262710958719254, -0.07785335928201675, -0.022374778985977173, -0.0817708671092987, 0.03487018868327141, -0.029547683894634247, -0.005801151506602764, -0.04905806481838226, 0.02166702039539814, -0.005992617458105087, -0.11015979200601578, -0.0308393482118845, 0.05095686763525009, -0.010311171412467957, 0.031786806881427765, -0.04617324098944664, -0.09157411754131317, 0.006420589052140713, -0.11849002540111542, 0.016280246898531914, -0.028407059609889984, -0.023433132097125053, -0.01440878864377737, 0.015474396757781506, -0.009521616622805595, 0.025167888030409813, 0.06421146541833878, -0.02171097695827484, -0.030941637232899666, -0.06835731118917465, 0.03074280172586441, -0.0014270101673901081, -0.004007979296147823, 0.003222087398171425, 0.0268495362251997, -0.06862882524728775, 0.07946494221687317, -0.013999230228364468, 0.03315926715731621, -0.007986468262970448, 0.03656792640686035, 0.009923783130943775, 0.0697653517127037, 0.05901985615491867, -0.01893511787056923, -0.08976630121469498, 0.08899316936731339, -0.08712100237607956, 0.04462917521595955, -0.08514243364334106, 0.020996108651161194, 0.030344098806381226, -0.039193253964185715, -0.00402164226397872, -0.046230658888816833, 0.036521490663290024, -0.005817084573209286, -0.012696418911218643, -0.07746623456478119, 0.07550348341464996, 0.04218454286456108, -0.010015941224992275, -0.002424961654469371, 0.08212081342935562, 0.05890953913331032, 0.023953698575496674, -0.001767810434103012, 0.030437715351581573, -0.06400702893733978, -0.03379131853580475, -0.09247767180204391, 0.11302445083856583, -0.01774044893682003, 0.004274298436939716, 0.08290484547615051, -0.02160218544304371, 0.006524897646158934, 0.06339986622333527, 0.08924776315689087, -0.09422831237316132, -0.010059158317744732, 0.09071087837219238, 0.05617143586277962, 0.08230188488960266, 0.0465853177011013, 0.01614946313202381, -0.05239386111497879, 0.017893586307764053, 0.057804301381111145, -0.01391006913036108, 0.06247490644454956, 0.036144766956567764, 0.04713922366499901, 0.023919865489006042, -0.0038782143965363503, 0.007813545875251293, -0.02174941822886467, -0.013273806311190128, -0.07201588898897171, -0.037237245589494705, 0.07555340975522995, -0.03177226334810257, 0.08040986210107803, 0.07055339962244034, -0.09208805859088898, 0.05642705783247948, 0.004612550139427185, -0.05031319707632065, 0.0010414078133180737, 0.098770372569561, -0.001608397695235908, 0.046751946210861206, -0.11464864760637283, 0.06781896948814392, -0.03997616469860077, 0.06479869037866592, -0.01006315741688013, -0.03510099649429321, -0.0036363841500133276, 0.06695260107517242, 0.08152879774570465, 0.020083345472812653, -0.010930563323199749, -0.03258099779486656, -0.09300528466701508, -0.027963219210505486, 0.08138058334589005, -0.012624758295714855, -0.08897827565670013, -0.024692652747035027, -0.012969769537448883, -0.01930255815386772, 0.017509700730443, 0.1616205871105194, -0.03731190040707588, -0.0514114648103714, -0.019485073164105415, 0.10054552555084229, -2.8669996027019806e-05, -0.08826442062854767, -0.0800241231918335, -0.059510037302970886, 0.04164430871605873, 0.013745177537202835, -0.05091278627514839, 0.02521866001188755, 0.08386386930942535, 0.06183091923594475, -0.026809783652424812, 0.02523311786353588, 0.013039928860962391, 0.03373462334275246, -0.07211276888847351, -0.03228246793150902, -0.02792643941938877, -0.02556450478732586, -0.036187123507261276, 0.019469473510980606, -0.0021493094973266125, -0.06982997804880142, 0.03882681578397751, -0.0019367545610293746, -0.019744519144296646, 0.022317824885249138, 0.010444262064993382, 0.07059932500123978, 0.025794727727770805, 0.035288408398628235, 0.03135279566049576, 0.0566762238740921, 0.11386867612600327, 0.003535778261721134, -0.03208962827920914, -0.025597339496016502, -0.011253892444074154, -0.02923794463276863, 0.09463649243116379, -0.05914255231618881, -0.030146606266498566, -0.008802848868072033, -0.021139204502105713, -0.02143789827823639, 0.042159970849752426, -0.003647301346063614, 0.02314841002225876, -0.04124969616532326, -0.04204202815890312, 0.05744973197579384, -0.09832766652107239, 0.058881811797618866, -0.03048650361597538, -0.05969224497675896, -0.0699595957994461, 0.03141476586461067, -0.012511608190834522, 0.05846767872571945, -0.011435642838478088, 0.024286707863211632, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.004008657298982143, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.2171594798564911, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.3737182319164276, 0, -0.2688813805580139, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.010957363061606884, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    print(embeddings)
    # 数据库参数
    from param_builder import MilvusParam
    db_param = MilvusParam.getUserParam()
    ans = milvusDatabase.searchEmbedding([embeddings], db_param)
    print(ans)
